#ifndef ___SIAMRPNTRACKER_H__
#define ___SIAMRPNTRACKER_H__

#include <math.h>
#include <stdio.h>
#include <stdlib.h>

#include <algorithm>
#include <fstream>
#include <iostream>
#include <opencv2/opencv.hpp>
#include <vector>
#include <xtensor/xadapt.hpp>
#include <xtensor/xarray.hpp>
#include <xtensor/xaxis_slice_iterator.hpp>
#include <xtensor/xbuilder.hpp>
#include <xtensor/xcomplex.hpp>
#include <xtensor/xcsv.hpp>
#include <xtensor/xhistogram.hpp>
#include <xtensor/xio.hpp>
#include <xtensor/xmasked_view.hpp>
#include <xtensor/xpad.hpp>
#include <xtensor/xshape.hpp>
#include <xtensor/xview.hpp>

#include "im2d.h"
#include "rga.h"
#include "rkmedia_buffer.h"
#include "rknn_api.h"
#include "timer_utils.h"

using namespace xt;
using namespace std;
using namespace cv;

#define M_PI 3.14159265358979323846

template <typename _Tp>
xarray<_Tp> softmax(xarray<_Tp> f, int dim) {
    xarray<_Tp> first = xt::exp(f);
    xarray<_Tp> second = xt::sum(first, {dim}, xt::keep_dims);
    return first / second;
}

template <typename _Tp>
xarray<_Tp> change(xarray<_Tp> r) {
    return xt::maximum(r, 1 / r);
}

template <typename _Tp>
xarray<_Tp> sz(xarray<_Tp> w, xarray<_Tp> h) {
    xarray<_Tp> pad = (w + h) * 0.5;
    return xt::sqrt((w + pad) * (h + pad));
}

template<typename _Tp>
_Tp change_t(_Tp r) {
    return std::max(r, 1 / r);
}

template<typename _Tp>
_Tp sz_t(_Tp w, _Tp h) {
    _Tp pad = (w + h) * 0.5;
    return std::sqrt((w + pad) * (h + pad));
}

static clock_t get_timestamp() {
    struct timeval tv;
    gettimeofday(&tv, NULL);

    return tv.tv_sec * 1000 + tv.tv_usec / 1000;
}

struct track_result {
    Rect bbox;
    float bscore;
    Point2d center_pos;
    float cost_time;
};

class SiamRpnTracker {
public:
    SiamRpnTracker();
    ~SiamRpnTracker();
    int preprocess_img(MEDIA_BUFFER fram_mb, cv::Size frame_size, Point2d K_center_pos);
    track_result track(cv::Size frame_size, bool redetect);
    int init(MEDIA_BUFFER fram_mb, cv::Size frame_size, Rect2d box);
    int initWithBGR(MEDIA_BUFFER fram_mb, cv::Size frame_size, Rect2d box);
    track_result postprocess_img(cv::Size frame_size);
    track_result postprocess_img_without_smooth(cv::Size frame_size);

protected:
    xarray<float> generate_anchors();
    xarray<float> convert_score(xarray<float> score);
    xarray<float> convert_box(xarray<float> delta, xarray<float> anchor);
    Rect2d bbox_clip(Rect2d box, Size boundary);

    int get_subwindowWithBGR(MEDIA_BUFFER fram_mb, cv::Size frame_size, Point2d center_pos, int model_size, int original_size, rknn_tensor_mem* tensor_mem);
    int get_subwindow(MEDIA_BUFFER fram_mb, cv::Size frame_size, Point2d center_pos, int model_size, int original_size, rknn_tensor_mem* tensor_mem);
    void printRKNNTensor(rknn_tensor_attr* attr);
    int printRKNNAttr(rknn_context ctx);

private:
    rknn_context template_session;
    rknn_context track_session;
    rknn_context rpn_session;

    void* template_modle;
    void* track_modle;
    void* rpn_modle;

    int track_instance_size;
    int track_exemplar_size;
    int track_base_size;
    double track_context_amount;
    float track_penalty_k;
    float track_window_influence;
    float track_lr;

    int anchor_stride;
    vector<float> anchor_ratios;
    vector<float> anchor_scales;
    int anchor_nump;

    int score_size;
    Point2d center_pos;
    Size2d box_size;

    Point2d usb_center_pos;
    Size2d usb_box_size;

    xarray<double> channel_average;
    xarray<float> anchor;
    xarray<double> window;

    bool has_inited;
    xarray<uchar> template_tensor;
    rknn_tensor_attr template_input_attrs;
    rknn_tensor_attr template_output_attrs;
    rknn_tensor_attr track_input_attrs;
    rknn_tensor_attr track_output_attrs;
    rknn_tensor_attr rpn_input_attrs[2];
    rknn_tensor_attr rpn_output_attrs[2];

    string tem_model_path = "/system/model/siamrpn_template_NHWC_128.rknn";
    string track_model_path = "/system/model/siamrpn_backbone_NHWC_288.rknn";
    string rpn_model_path = "/system/model/siamrpn_rpn.rknn"; 

    size_t buffer_output_cls_size = 10 * 21 * 21;
    size_t buffer_output_bbox_size = 20 * 21 * 21;
    size_t buffer_output_number = 5 * 21 * 21;

    rknn_tensor_mem* tem_inputs_mem;
    rknn_tensor_mem* tem_outputs_mem;
    rknn_tensor_mem* track_inputs_mem;
    rknn_tensor_mem* track_outputs_mem;
    rknn_tensor_mem* rpn_outputs_mem[2];

    float out_scales[2];
    uint32_t out_zps[2];
};

#endif
